import os
import os.path as osp
import shutil
from sklearn.model_selection import train_test_split


def create_train_val_test_splits(output_base_dir, train_ratio=0.7, val_ratio=0.2, test_ratio=0.1):
    """在输出目录中划分训练集、验证集和测试集"""
    a_dir = os.path.join(output_base_dir, 'A')
    files = [f for f in os.listdir(a_dir) if f.endswith('.jpg')]

    # 划分数据集
    train_files, remaining = train_test_split(
        files, test_size=(1 - train_ratio), random_state=42
    )
    val_ratio_adjusted = val_ratio / (val_ratio + test_ratio)
    val_files, test_files = train_test_split(
        remaining, test_size=1 - val_ratio_adjusted, random_state=42
    )

    splits = {'train': train_files, 'val': val_files, 'test': test_files}
    for split_name, file_list in splits.items():
        split_path = os.path.join(output_base_dir, split_name)
        os.makedirs(os.path.join(split_path, 'A'), exist_ok=True)
        os.makedirs(os.path.join(split_path, 'B'), exist_ok=True)
        os.makedirs(os.path.join(split_path, 'label'), exist_ok=True)
        for filename in file_list:
            src_a = os.path.join(output_base_dir, 'A', filename)
            src_b = os.path.join(output_base_dir, 'B', filename)
            src_label = os.path.join(output_base_dir, 'label', filename)

            dst_a = os.path.join(split_path, 'A', filename)
            dst_b = os.path.join(split_path, 'B', filename)
            dst_label = os.path.join(split_path, 'label', filename)

            if os.path.exists(src_a):
                shutil.copy(src_a, dst_a)
            if os.path.exists(src_b):
                shutil.copy(src_b, dst_b)
            if os.path.exists(src_label):
                shutil.copy(src_label, dst_label)
    print("已完成训练集、验证集、测试集划分。")


def generate_txt_from_dir(src_dir, dst_dir, split):
    """Generate .txt file for LEVIR-CD dataset.
    Args:
        src_dir (str): path of the source dataset.
        dst_dir (str): Path to save .txt file.
        split (str): sub_dirs. 'train', 'val' or 'test'
    """
    src_dir = osp.join(src_dir, split)
    sub_dir_1 = osp.join(src_dir, 'A')
    sub_dir_2 = osp.join(src_dir, 'B')
    ann_dir = osp.join(src_dir, 'label')

    file_list = []
    for img_name in sorted(os.listdir(ann_dir)):
        assert osp.exists(osp.join(sub_dir_1, img_name)) and \
               osp.exists(osp.join(sub_dir_2, img_name)), \
            f'{img_name} is not in {sub_dir_1} or {sub_dir_2}'

        file_list.append([
            os.path.splitext(img_name)[0]
        ])

    with open('{}.txt'.format(osp.join(dst_dir, split)), 'w') as f:
        for item in file_list:
            f.write(' '.join(item) + '\n')


def create_directory_if_not_exists(specified_root):
    if not os.path.exists(specified_root):
        try:
            os.makedirs(specified_root)
            print("'{specified_root}' created ....")
        except OSError as error:
            print("Failed to create directory '{specified_root}'. ")
    else:
        # print(f"Directory '{specified_root}' already exists.")
        pass


if __name__ == '__main__':

    # 1）划分数据集
    output_base_dir = r'D:\0_Data\_cd_data\2_gen_data\ng_rgb'  # 你可以修改为你想要的输出路径
    create_train_val_test_splits(output_base_dir, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1)

    # 2）创建list文件夹，并在这三个文件夹内存储对应的txt文件
    print('Making .txt files ...')
    list_dir = r"D:\0_Data\_cd_data\2_gen_data\ng_rgb\list"
    create_directory_if_not_exists(list_dir)
    generate_txt_from_dir(output_base_dir, list_dir, 'train')
    generate_txt_from_dir(output_base_dir, list_dir, 'val')
    generate_txt_from_dir(output_base_dir, list_dir, 'test')
    print("数据预处理完成！")